
import logging
from pinecone import Pinecone
from sklearn.datasets import load_digits
from sklearn.model_selection import train_test_split
import numpy as np
from collections import Counter
from tqdm import tqdm

# 配置 logging，包含日期
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s',
    datefmt='%Y-%m-%d %H:%M:%S'
)

# 假设 pinecone_example.py 已经创建并连接了索引，这里只需连接即可
pinecone = Pinecone(api_key="pcsk_3Du3r4_PPKjDJ3RXs6svx1MWTwTgVEK3VSpdFekSf1aH5TqmFqzzYoSUUCxhRU2ZRXQ1jQ")
index_name = "mnist-index"
index = pinecone.Index(index_name)
logging.info(f"已连接到索引 '{index_name}'。")

# 加载 MNIST 数据集
digits = load_digits(n_class=10)
X = digits.data
y = digits.target

# 划分训练集和测试集（80%训练，20%测试）
X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.2, random_state=42, stratify=y
)

# 组装 Pinecone 向量格式
vectors = []
for i in range(len(X_train)):
    vector_id = str(i)
    vector_values = X_train[i].tolist()
    metadata = {"label": int(y_train[i])}
    vectors.append((vector_id, vector_values, metadata))


# 上传数据到 Pinecone，带进度条
batch_size = 1000
logging.info(f"开始上传训练数据到 Pinecone，共 {len(vectors)} 条...")
for i in tqdm(range(0, len(vectors), batch_size), desc="上传进度", ncols=80):
    batch = vectors[i:i + batch_size]
    index.upsert(batch)
logging.info(f"成功上传 {len(vectors)} 条数据到 Pinecone。")

# 用测试集评估准确率，k=11，带进度条
correct = 0
logging.info(f"开始用测试集评估 Pinecone 索引，k=11...")
for i in tqdm(range(len(X_test)), desc="测试进度", ncols=80):
    query_data = X_test[i].tolist()
    results = index.query(
        vector=query_data,
        top_k=11,
        include_metadata=True
    )
    labels = [match['metadata']['label'] for match in results['matches']]
    if labels:
        pred = Counter(labels).most_common(1)[0][0]
        if pred == y_test[i]:
            correct += 1
accuracy = correct / len(X_test)
logging.info(f"k=11 时的准确率为：{accuracy:.4f}")
